# Copyright Dechin CHEN 2025

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from functools import wraps
import mindspore as ms
from mindspore import ops
from mindspore import numpy as msnp


int32 = ms.int32
int64 = ms.int64
float16 = ms.float16
float32 = ms.float32
float64 = ms.float64
bool_ = ms.bool_
Tensor = ms.Tensor
tensor = ms.Tensor
long = ms.int64
_Metric = ms.train.Metric
stop_gradient = ops.stop_gradient
depend = ops.Depend()
tensor = ms.Tensor
Tensor = ms.Tensor
from_numpy = ms.from_numpy


_MINDSPORE_OPS = {}

def register_op(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    _MINDSPORE_OPS[func.__name__] = wrapper
    return wrapper


@register_op
def get_parameters(model):
    return model.trainable_params()


@register_op
def Sum(ipt, *args, **kwargs):
    return msnp.sum(ipt, *args, **kwargs)


@register_op
def where(ipt, *args, **kwargs):
    return msnp.where(ipt, *args, **kwargs)

pad = ms.mint.nn.functional.pad

@register_op
def reduce_sum(arr, dim):
    return ops.ReduceSum(keep_dims=False)(arr, dim)


@register_op
def reduce_max(input, dim=None, keepdim=False):
    return ops.ReduceMax(keep_dims=False)(arr, dim)


@register_op
def reduce_min(input, dim=None, keepdim=False):
    return ops.ReduceMin(keep_dims=False)(arr, dim)

@register_op
def is_tensor(input, *args, **kwargs):
    return ops.is_tensor(input, *args, **kwargs)

@register_op
def is_complex(input, *args, **kwargs):
    return ops.is_complex(input, *args, **kwargs)

@register_op
def is_nonzero(input, *args, **kwargs):
    return ops.is_nonzero(input, *args, **kwargs)

@register_op
def numel(input, *args, **kwargs):
    return ops.numel(input, *args, **kwargs)

@register_op
def to_numpy(arr, *args, **kwargs):
    return arr.asnumpy()

@register_op
def asarray(input, *args, **kwargs):
    return ms.numpy.asarray(input, *args, **kwargs)

@register_op
def zeros(input, *args, **kwargs):
    return ms.mint.zeros(input, *args, **kwargs)

@register_op
def zeros_like(input, *args, **kwargs):
    return ms.mint.zeros_like(input, *args, **kwargs)


@register_op
def ones(input, *args, **kwargs):
    return ms.mint.ones(input, *args, **kwargs)


@register_op
def ones_like(input, *args, **kwargs):
    return ms.mint.ones_like(input, *args, **kwargs)
